package com.test.limit;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.RedisZSetCommands;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 时间窗口限流
 *
 * @author zxs
 * @since 2020/06/23
 */
@Service("timeWindowLimit")
public class TimeWindowLimit implements Limiter {

    private static final String prefix = "TWL:";

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 每X秒允许行为发生Y次
     *
     * @param key      行为
     * @param period   时间单位
     * @param maxCount 最大次数
     * @return true/false
     */
    public boolean isAllow(String key, Integer maxCount, Integer period) {
        ZSetOperations<String,Object> zSetOperations = redisTemplate.opsForZSet();
        long now = System.currentTimeMillis();
        String fullKey = prefix + key;
        //删除时间窗口外的数据
        zSetOperations.removeRangeByScore(fullKey, 0d, now - (period * 1000));
        //获取当前窗口内总数量
        Long count = zSetOperations.zCard(fullKey);
        //大于最大值，拒绝
        if (count != null && count > maxCount) {
            return false;
        }
        //添加一条记录
        zSetOperations.add(fullKey, now, now);
        redisTemplate.expire(fullKey,period+1, TimeUnit.SECONDS);
        return true;
    }

}
